import os
import sqlite3
import json
import platform
import shutil
import time
import subprocess
import argparse
from pathlib import Path

class CursorAccountSwitcher:
    def __init__(self):
        self.system = platform.system()
        self.base_path = self._get_base_path()
        self.db_path = self._get_db_path()
        self.storage_path = self._get_storage_path()
        self.main_js_path = self._get_main_js_path()

    def _get_base_path(self):
        """获取Cursor基础路径"""
        if self.system == "Darwin":  # macOS
            return os.path.join(
                os.path.expanduser("~"),
                "Library/Application Support/Cursor"
            )
        elif self.system == "Windows":
            return os.path.join(os.getenv("APPDATA"), "Cursor")
        else:  # Linux
            return os.path.join(os.path.expanduser("~"), ".config", "Cursor")

    def _get_db_path(self):
        """获取数据库路径"""
        if self.system == "Darwin":
            paths = [
                os.path.join(self.base_path, "User/globalStorage/state.vscdb"),
                os.path.join(self.base_path, "Local Storage/leveldb/state.vscdb"),
                os.path.join(self.base_path, "User/state.vscdb")
            ]
            for path in paths:
                if os.path.exists(path):
                    return path
            return paths[0]
        elif self.system == "Linux":
            paths = [
                os.path.join(self.base_path, "User/globalStorage/state.vscdb"),
                os.path.join(self.base_path, "Local Storage/leveldb/state.vscdb"),
                os.path.join(self.base_path, "User/state.vscdb"),
                os.path.join(self.base_path, "Service Worker/Database/state.vscdb")
            ]
            for path in paths:
                if os.path.exists(path):
                    return path
            return paths[0]
        else:
            return os.path.join(self.base_path, "User", "globalStorage", "state.vscdb")

    def _get_storage_path(self):
        """获取storage.json路径"""
        if self.system == "Darwin":
            paths = [
                os.path.join(self.base_path, "User/globalStorage/storage.json"),
                os.path.join(self.base_path, "Local Storage/leveldb/storage.json"),
                os.path.join(self.base_path, "User/storage.json")
            ]
            for path in paths:
                if os.path.exists(path):
                    return path
            return paths[0]
        elif self.system == "Linux":
            paths = [
                os.path.join(self.base_path, "User/globalStorage/storage.json"),
                os.path.join(self.base_path, "Local Storage/leveldb/storage.json"),
                os.path.join(self.base_path, "User/storage.json"),
                os.path.join(self.base_path, "Service Worker/Database/storage.json")
            ]
            for path in paths:
                if os.path.exists(path):
                    return path
            return paths[0]
        else:
            return os.path.join(self.base_path, "User", "globalStorage", "storage.json")

       
    def _get_main_js_path(self):
        """获取main.js路径"""
        if self.system == "Darwin":
            paths = [
                "/Applications/Cursor.app/Contents/Resources/app/out/vs/code/electron-main/main.js",
                "/Applications/Cursor.app/Contents/Resources/app/main.js"
            ]
            for path in paths:
                if os.path.exists(path):
                    return path
            return paths[0]
        elif self.system == "Linux":
            paths = [
                "/usr/share/cursor/resources/app/out/vs/code/electron-main/main.js",
                "/opt/cursor/resources/app/out/vs/code/electron-main/main.js",
                os.path.expanduser("~/.local/share/cursor/resources/app/out/vs/code/electron-main/main.js")
            ]
            for path in paths:
                if os.path.exists(path):
                    return path
            return paths[0]
        else:
            local_app_data = os.getenv("LOCALAPPDATA")
 
            return os.path.join(local_app_data, "Programs", "cursor","resources", "app", "out", "main.js")

    def _kill_cursor(self):
        """关闭Cursor进程"""
        try:
            if self.system == "Darwin":
                subprocess.run(["pkill", "-f", "Cursor.app"], check=False)
                subprocess.run(["killall", "Cursor"], check=False)
            elif self.system == "Linux":
                subprocess.run(["pkill", "-f", "cursor"], check=False)
                subprocess.run(["killall", "cursor"], check=False)
            else:
                subprocess.run(["taskkill", "/F", "/IM", "Cursor.exe"], check=False)
            
            time.sleep(1)
            return True
        except Exception as e:
            print(f"注意: 未找到Cursor进程，继续执行...")
            return True

    def _set_file_permissions(self, file_path):
        """设置文件权限"""
        if self.system in ["Darwin", "Linux"]:
            try:
                subprocess.run(["chmod", "666", file_path], check=False)
                return True
            except Exception as e:
                print(f"警告: 无法修改文件权限: {file_path}")
                return False
        return True

    def _backup_file(self, file_path):
        """备份文件"""
        if os.path.exists(file_path):
            backup_path = f"{file_path}.backup"
            try:
                shutil.copy2(file_path, backup_path)
                print(f"已创建备份: {backup_path}")
                return True
            except Exception as e:
                print(f"创建备份失败: {str(e)}")
                return False
        return False

    def _update_db(self, email, token):
        """更新数据库中的登录凭证"""
        if not os.path.exists(self.db_path):
            print(f"错误: 找不到数据库文件 {self.db_path}")
            return False

        # 设置文件权限
        if self.system in ["Darwin", "Linux"]:
            self._set_file_permissions(self.db_path)

        try:
            conn = sqlite3.connect(self.db_path)
            cursor = conn.cursor()

            # 检查表是否存在
            cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='ItemTable'")
            if not cursor.fetchone():
                print("警告: 未找到ItemTable表，尝试其他表名...")
                cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='item'")
                if cursor.fetchone():
                    table_name = "item"
                else:
                    print("错误: 找不到合适的表")
                    return False
            else:
                table_name = "ItemTable"

            # 处理token格式
            processed_token = token.split("%3A%3A")[-1] if "%3A%3A" in token else token

            # 更新登录凭证
            updates = [
                ("cursor.email", email),
                ("cursor.accessToken", processed_token),
                ("cursorAuth/refreshToken", processed_token),
                ("cursorAuth/accessToken", processed_token),
                ("cursorAuth/cachedEmail", email)
            ]

            for key, value in updates:
                cursor.execute(
                    f"INSERT OR REPLACE INTO {table_name} (key, value) VALUES (?, ?)",
                    (key, value)
                )

            conn.commit()
            conn.close()
            print("数据库更新成功")
            return True
        except Exception as e:
            print(f"更新数据库失败: {str(e)}")
            return False

    def switch_account(self, email, token):
        """切换账号"""
        print(f"开始切换账号... (系统: {self.system})")
        print(f"邮箱: {email}")
        print(f"数据库路径: {self.db_path}")
        
        # 1. 尝试关闭Cursor
        print("正在检查Cursor进程...")
        self._kill_cursor()  # 无论结果如何都继续
        
        # 2. 备份数据库
        print("正在备份数据库...")
        if os.path.exists(self.db_path):
            self._backup_file(self.db_path)
        else:
            print(f"警告: 数据库文件不存在: {self.db_path}")
            return False

        # 3. 更新数据库
        print("正在更新数据库...")
        if not self._update_db(email, token):
            return False

        print("账号切换完成！")
        return True

    def get_current_account(self):
        """获取当前Cursor账号信息，返回邮箱和token"""
        print(f"正在获取当前Cursor账号信息... (系统: {self.system})")
        print(f"数据库文件 {self.db_path}")
        # 检查数据库文件是否存在
        if not os.path.exists(self.db_path):
            print(f"错误: 找不到数据库文件 {self.db_path}")
            return None, None
        
        # 设置文件权限
        if self.system in ["Darwin", "Linux"]:
            self._set_file_permissions(self.db_path)
        
        try:
            conn = sqlite3.connect(self.db_path)
            cursor = conn.cursor()
            
            # 检查表是否存在
            cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='ItemTable'")
            if not cursor.fetchone():
                print("警告: 未找到ItemTable表，尝试其他表名...")
                cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='item'")
                if cursor.fetchone():
                    table_name = "item"
                else:
                    print("错误: 找不到合适的表")
                    return None, None
            else:
                table_name = "ItemTable"
            
            # 查询email
            cursor.execute(f"SELECT value FROM {table_name} WHERE key = 'cursorAuth/cachedEmail'")
            email_result = cursor.fetchone()
            email = email_result[0] if email_result else None
            
            # 查询token
            cursor.execute(f"SELECT value FROM {table_name} WHERE key = 'cursorAuth/accessToken'")
            token_result = cursor.fetchone()
            token = token_result[0] if token_result else None
            
            conn.close()
            
            if email:
                print(f"获取到当前账号: {email}")
                return email, token
            else:
                print("未找到登录信息")
                return None, None
        except Exception as e:
            print(f"获取当前账号信息失败: {str(e)}")
            return None, None

    def modify_machine_code(self):
        """修改Cursor机器码（包括修改main.js和storage.json）"""
        print("开始修改机器码...")
        success_count = 0
        
        # 1. 修改storage.json
        if os.path.exists(self.storage_path):
            print(f"正在处理storage.json: {self.storage_path}")
            try:
                # 创建备份
                self._backup_file(self.storage_path)
                
                # 读取storage.json
                with open(self.storage_path, 'r', encoding='utf-8') as f:
                    try:
                        storage_data = json.load(f)
                    except json.JSONDecodeError:
                        print("错误: storage.json不是有效的JSON格式")
                        storage_data = {}
                
                # 生成新的随机ID
                import uuid
                new_ids = {
                    "telemetry.devDeviceId": str(uuid.uuid4()),
                    "telemetry.macMachineId": str(uuid.uuid4()),
                    "telemetry.machineId": str(uuid.uuid4()),
                    "telemetry.sqmId": str(uuid.uuid4())
                }
                
                # 更新机器码
                print("正在更新机器码标识...")
                for key, value in new_ids.items():
                    storage_data[key] = value
                    
                # 写入修改后的内容
                with open(self.storage_path, 'w', encoding='utf-8') as f:
                    json.dump(storage_data, f, indent=2)
                    
                print("storage.json更新成功")
                success_count += 1
            except Exception as e:
                print(f"修改storage.json失败: {str(e)}")
        else:
            print(f"警告: 找不到storage.json文件: {self.storage_path}")
        
        # 2. 修改main.js
        if os.path.exists(self.main_js_path):
            print(f"正在处理main.js: {self.main_js_path}")
            #处理main.js核心逻辑

            #正则匹配找到getMacMachineId() { return this.c.macMachineId } 然后改为
            #getMacMachineId() { return this.c.macMachineId+new Date().getTime() }


            #正则匹配找到getMachineId() { return this.c.machineId } 改为
            #getMachineId() { return this.c.machineId +new Date().getTime() }

            # 设置文件权限
            if self.system in ["Darwin", "Linux"]:
                self._set_file_permissions(self.main_js_path)
                
            try:
                # 读取main.js文件
                with open(self.main_js_path, 'r', encoding='utf-8') as f:
                    content = f.read()
                    
                # 创建备份
                backup_path = f"{self.main_js_path}.backup"
                if not os.path.exists(backup_path):
                    with open(backup_path, 'w', encoding='utf-8') as f:
                        f.write(content)
                    print(f"已创建备份: {backup_path}")
                    
                # 使用正则表达式进行替换
                import re
                
                # 新的匹配模式：匹配简洁版函数定义
                machine_id_patterns = [
                    # 匹配原始形式的函数定义
                    r'async\s+(\w+)\s*\(\)\s*\{\s*return\s+this\.[\w.]+(?:\?\?|\?)\s*this\.([\w.]+)\.machineId\s*\}',
                    # 匹配简洁形式函数定义 getMachineId() { return this.c.machineId }
                    r'(\w+MachineId)\s*\(\)\s*\{\s*return\s+this\.([\w.]+)\.machineId\s*\}'
                ]
                
                # 匹配简洁形式的mac机器码函数定义
                mac_machine_id_patterns = [
                    # 匹配原始形式的函数定义
                    r'async\s+(\w+)\s*\(\)\s*\{\s*return\s+this\.[\w.]+(?:\?\?|\?)\s*this\.([\w.]+)\.macMachineId\s*\}',
                    # 匹配简洁形式函数定义 getMacMachineId() { return this.c.macMachineId }
                    r'(getMacMachineId)\s*\(\)\s*\{\s*return\s+this\.([\w.]+)\.macMachineId\s*\}'
                ]
                
                # 查找匹配的数量
                machine_id_matches = sum(len(re.findall(pattern, content)) for pattern in machine_id_patterns)
                mac_machine_id_matches = sum(len(re.findall(pattern, content)) for pattern in mac_machine_id_patterns)
                
                if machine_id_matches == 0 and mac_machine_id_matches == 0:
                    print("警告: 未找到需要修改的机器码验证函数")
                else:
                    print(f"找到 {machine_id_matches} 个machineId函数和 {mac_machine_id_matches} 个macMachineId函数")
                    
                    # 替换machineId函数 - 添加时间戳
                    modified_content = content
                    for pattern in machine_id_patterns:
                        modified_content = re.sub(
                            pattern,
                            lambda m: f"{m.group(1)}() {{ return this.{m.group(2)}.machineId+new Date().getTime() }}",
                            modified_content
                        )
                    
                    # 替换macMachineId函数 - 添加时间戳
                    for pattern in mac_machine_id_patterns:
                        modified_content = re.sub(
                            pattern,
                            lambda m: f"{m.group(1)}() {{ return this.{m.group(2)}.macMachineId+new Date().getTime() }}",
                            modified_content
                        )
                    
                    # 检查是否有实际修改
                    if content == modified_content:
                        print("警告: main.js未修改任何内容")
                    else:
                        # 写入修改后的内容
                        with open(self.main_js_path, 'w', encoding='utf-8') as f:
                            f.write(modified_content)
                        print("main.js更新成功")
                        success_count += 1
                
            except Exception as e:
                print(f"修改main.js失败: {str(e)}")
        else:
            print(f"警告: 找不到main.js文件: {self.main_js_path}")
        
        # 汇总结果
        if success_count == 0:
            print("机器码修改失败: 所有文件操作都未成功")
            return False
        elif success_count == 2:
            print("机器码修改完全成功: storage.json和main.js均已更新")
            return True
        else:
            print(f"机器码修改部分成功: 更新了{success_count}/2个文件")
            return True

def main():
    # 创建命令行参数解析器
    parser = argparse.ArgumentParser(description='Cursor账号切换工具')
    parser.add_argument('-e', '--email', required=True, help='新的邮箱地址')
    parser.add_argument('-t', '--token', required=True, help='新的token')
    parser.add_argument('--force', action='store_true', help='强制执行，不询问确认')
    
    # 解析命令行参数
    args = parser.parse_args()
    
    # 创建切换器实例
    switcher = CursorAccountSwitcher()
    
    # 如果不是强制执行，显示确认信息
    if not args.force:
        print(f"即将切换到以下账号：")
        print(f"邮箱: {args.email}")
        print(f"Token: {args.token[:10]}...{args.token[-10:] if len(args.token) > 20 else args.token}")
        confirm = input("确认继续？(y/N): ")
        if confirm.lower() != 'y':
            print("操作已取消")
            return
    
    # 执行账号切换
    if switcher.switch_account(args.email, args.token):
        print("账号切换成功！")
        print("请重新启动Cursor。")
    else:
        print("账号切换失败！")
        exit(1)

if __name__ == "__main__":
    main()